from __future__ import print_function
import argparse
import os
import csv
import numpy as np
import random
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from data_utils.data_util import PointcloudScaleAndTranslate
from data_utils.ModelNetDataLoader import ModelNetDataLoader
from data_utils.ModelNetDataLoader_concate import ModelNetDataLoaderHybrid

from models.pointnet import PointNetCls, feature_transform_regularizer
from models.pointnet2 import PointNet2ClsMsg
from models.dgcnn import DGCNN
from models.pointcnn import PointCNNCls
from models.lowFreModel_random import LowFreqEnhancedModel  # 新增模型导入

from utils import progress_bar, log_row
import sys
sys.path.append("./emd/")
import emd_module as emd

def gen_train_log(args):
    if not os.path.isdir('logs_train'):
        os.mkdir('logs_train')
    logname = ('logs_train/%s_%s_%s.csv' % (args.data, args.model, args.name))

    if os.path.exists(logname):
        with open(logname, 'a') as logfile:
            log_row(logname, [''])
            log_row(logname, [''])

    with open(logname, 'a') as logfile:
        logwriter = csv.writer(logfile, delimiter=',')
        logwriter.writerow(['model type', 'data set', 'seed', 'train batch size',
                            'number of points in one batch', 'number of epochs', 'optimizer',
                            'learning rate', 'resume checkpoint path',
                            'feature transform', 'lambda for feature transform regularizer', 'data augment'])
        logwriter.writerow([args.model, args.data, args.seed, args.batch_size, args.num_points,
                            args.epochs, args.optimizer, args.lr, args.resume,
                            args.feature_transform, args.lambda_ft, args.augment])
        logwriter.writerow(['Note', args.note])
        logwriter.writerow([''])


def save_ckpt(args, epoch, model, optimizer, acc_list):
    if not os.path.isdir('checkpoints'):
        os.mkdir('checkpoints')
    if not os.path.isdir('checkpoints/%s_%s_%s' % (args.data, args.model, args.name)):
        os.mkdir('checkpoints/%s_%s_%s' % (args.data, args.model, args.name))
    if acc_list[-1] > max(acc_list[:-1]):
        print('=====> Saving checkpoint...')
        print('the best test acc is', acc_list[-1])
        state = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'args': args,
            'acc_list': acc_list,
        }
        torch.save(state, 'checkpoints/%s_%s_%s/best.pth' % (args.data, args.model, args.name))
        print('Successfully save checkpoint at epoch %d' % epoch)


def cal_loss(pred, gold, smoothing=True):
    ''' Calculate cross entropy loss, apply label smoothing if needed. '''
    gold = gold.contiguous().view(-1)

    if smoothing:
        eps = 0.2
        n_class = pred.size(1)

        one_hot = torch.zeros_like(pred).scatter(1, gold.view(-1, 1), 1)
        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
        log_prb = F.log_softmax(pred, dim=1)

        loss = -(one_hot * log_prb).sum(dim=1).mean()
    else:
        loss = F.cross_entropy(pred, gold, reduction='mean')

    return loss


def test(model, test_loader, criterion):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0
    with torch.no_grad():
        for j, data in enumerate(test_loader, 0):
            points, label = data
            points, label = points.to(device), label.to(device)
            
            # 生成多尺度特征(示例用随机生成)
            original_feats = [
                torch.randn(points.size(0), 1024, 256).to(device),
                torch.randn(points.size(0), 512, 512).to(device),
                torch.randn(points.size(0), 256, 1024).to(device)
            ]
            
            # 前向传播
            if 'lowfreq' in args.model:
                pred = model(points[:, :, :3], original_feats)  # 仅使用坐标
            else:
                points = points.transpose(2, 1)
                pred, _ = model(points)
            
            loss = criterion(pred, label.long())
            test_loss += loss.item()
            
            pred_choice = pred.data.max(1)[1]
            correct += pred_choice.eq(label.data).cpu().sum()
            total += label.size(0)
            progress_bar(j, len(test_loader), 'Test Loss: %.3f | Test Acc: %.3f%% (%d/%d)'
                        % (test_loss/(j+1), 100.*correct.item()/total, correct, total))
    return test_loss/len(test_loader), 100.*correct.item()/total


if __name__ == '__main__':
    ########################################
    ## Set hypeparameters
    ########################################
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='choose model type:[pointnet|pointnet2|dgcnn|lowfreq]', help='choose model type')
    parser.add_argument('--data', type=str, default='modelnet40', help='choose data set')
    parser.add_argument('--seed', type=int, default=0, help='manual random seed')
    parser.add_argument('--batch_size', type=int, default=8, help='input batch size')
    parser.add_argument('--num_points', type=int, default=1024, help='input batch size')
    parser.add_argument('--epochs', type=int, default=300, help='number of epochs to train for')
    parser.add_argument('--optimizer', type=str, default='Adam', help='optimizer for training')
    parser.add_argument('--lr', default=0.001, type=float, help='learning rate in training')
    parser.add_argument('--resume', type=str, default='/', help='resume path')
    parser.add_argument('--feature_transform', type=int, default=1, help="use feature transform")
    parser.add_argument('--lambda_ft', type=float, default=0.001, help="lambda for feature transform")
    parser.add_argument('--augment', type=int, default=1, help='data argment to increase robustness')
    parser.add_argument('--name', type=str, default='train', help='name of the experiment')
    parser.add_argument('--note', type=str, default='', help='notation of the experiment')
    parser.add_argument('--normal', action='store_true', default=False,
                        help='Whether to use normal information [default: False]')
    parser.add_argument('--beta', default=1, type=float, help='hyperparameter beta')
    parser.add_argument('--cutmix_prob', default=0.5, type=float, help='cutmix probability')
    args = parser.parse_args()
    args.feature_transform, args.augment = bool(args.feature_transform), bool(args.augment)
    ### Set random seed
    args.seed = args.seed if args.seed > 0 else random.randint(1, 10000)

    # dataset path
    DATA_PATH = '/media/shangli211/4TB_SSD/program_file/Data/'
    Def_PATH = '/media/shangli211/4TB_SSD/program_file/Data/def_data'


    ########################################
    ## Intiate model
    ########################################
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_classes = 40

    if args.model == 'dgcnn_kcutmix':
        model = DGCNN(num_classes)
        model = model.to(device)
        model = nn.DataParallel(model)
        optimizer = torch.optim.SGD(model.parameters(), lr=args.lr * 100,
                                    momentum=0.9, weight_decay=1e-4)
        scheduler_c = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 250,
                                                                 eta_min=1e-3)
    else:
        if args.model == 'pointnet_kcutmix':
            model = PointNetCls(num_classes, args.feature_transform)
            model = model.to(device)
        elif args.model == 'pointnet2_kcutmix':
            model = PointNet2ClsMsg(num_classes)
            model = model.to(device)
            model = nn.DataParallel(model)

        elif args.model == 'rscnn_kcutmix':
            from models.rscnn import RSCNN
            import models.rscnn_utils.pointnet2_utils as pointnet2_utils
            model = RSCNN(num_classes)
            model = model.to(device)
            model = nn.DataParallel(model)

        elif args.model == "lowfreq":
            from models.lowFreModel_random import LowFreqEnhancedModel
            model = LowFreqEnhancedModel(
                lmax=16,
                sigma=0.1,
                pc_size=args.num_points,
                device=device
            ).to(device)

        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=args.lr,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=1e-4
        )
        scheduler_c = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
    if len(args.resume) > 1:
        print('=====> Loading from checkpoint...')
        checkpoint = torch.load('./checkpoints/%s.pth' % args.resume)
        args = checkpoint['args']

        torch.manual_seed(args.seed)
        print("Random Seed: ", args.seed)

        """if args.optimizer == 'SGD':
            optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
        elif args.optimizer == 'Adam':
            optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999))"""

        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        START_EPOCH = checkpoint['epoch'] + 1
        acc_list = checkpoint['acc_list']
        if args.model == 'dgcnn_kcutmix':
            scheduler_c = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 250, eta_min=1e-3)
        else:
            scheduler_c = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
        print('Successfully resumed!')

    else:
        print('=====> Building new model...')
        torch.manual_seed(args.seed)
        print("Random Seed: ", args.seed)


        START_EPOCH = 0
        acc_list = [0]

        print('Successfully built!')

    ########################################
    ## Load data
    ########################################
    print('======> Loading data')

    # TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_points, split='train',
    #                                     normal_channel=args.normal)
    TRAIN_DATASET = ModelNetDataLoaderHybrid(root=DATA_PATH,def_root=Def_PATH,npoint=1024,split='train',normal_channel=args.normal)

    TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_points, split='test',
                                        normal_channel=args.normal)
    train_loader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True,
                                               num_workers=4, drop_last=True)
    test_loader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False,
                                              num_workers=4, drop_last=False)

    PointcloudScaleAndTranslate = PointcloudScaleAndTranslate()
    print('======> Successfully loaded!')

    gen_train_log(args)
    logname = ('logs_train/%s_%s_%s.csv' % (args.data, args.model, args.name))

    ########################################
    ## Train
    ########################################
if args.model == 'dgcnn_kcutmix':
    criterion = cal_loss
else:
    criterion = F.cross_entropy  # nn.CrossEntropyLoss()

if args.resume == '/':
    log_row(logname, ['Epoch', 'Train Loss', 'Train Acc', 'Test Loss', 'Test Acc', 'learning Rate'])

for epoch in range(START_EPOCH, args.epochs):
    print('\nEpoch: %d' % epoch)
    scheduler_c.step(epoch)
    model.train()

    correct = 0
    total = 0
    for i, data in enumerate(train_loader, 0):
        points, target = data
        points, target = points.to(device), target.to(device)
        points = PointcloudScaleAndTranslate(points)

        if args.model == 'rscnn_kcutmix':
            fps_idx = pointnet2_utils.furthest_point_sample(points, args.num_points)  # (B, npoint)
            fps_idx = fps_idx[:, np.random.choice(args.num_points, args.num_points, False)]
            points = pointnet2_utils.gather_operation(points.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous()  # (B, N, 3)

        # cutmix
        optimizer.zero_grad()
        r = np.random.rand(1)

        # 先保存原始的 points 用于后续处理
        original_points = points.clone()

        points = points.transpose(2, 1)

        # 生成 original_feats
        original_feats = [
            torch.randn(points.size(0), 1024, 256).to(device),
            torch.randn(points.size(0), 512, 512).to(device),
            torch.randn(points.size(0), 256, 1024).to(device)
        ]

        if args.beta > 0 and r < args.cutmix_prob:
            lam = np.random.beta(args.beta, args.beta)
            B = points.size()[0]

            rand_index = torch.randperm(B).cuda()
            target_a = target
            target_b = target[rand_index]

            # 统一点云尺寸到1024
            def pad_to_1024(pc):
                B, N, C = pc.shape
                assert C == 3, f"Input tensor should have 3 channels (coordinates), but got {C}"
                if N != 1024:
                    return F.pad(pc, (0, 0, 0, 1024 - N, 0, 0))  # 填充到固定1024点
                return pc

            # 处理所有点云张量
            point_a = pad_to_1024(original_points)
            point_b = pad_to_1024(original_points[rand_index])
            point_c = pad_to_1024(original_points[rand_index])

            # 转换到GPU
            point_a, point_b, point_c = point_a.to(device), point_b.to(device), point_c.to(device)

            # 计算EMD匹配
            remd = emd.emdModule().cuda()
            dis, ind = remd(point_a, point_b, 0.005, 300)  # ind形状应为[B, 1024]

            # 直接应用索引（已确保尺寸一致）
            point_c = point_c.gather(1, ind.long().unsqueeze(-1).expand(-1, -1, 3))

            # 生成混合区域
            int_lam = int(1024 * lam)  # 直接使用1024作为基数
            int_lam = max(1, min(int_lam, 1023))  # 限制在1-1023之间

            # 随机选择中心点
            center_point = torch.randint(0, 1024, (B,), device=device)

            # 创建查询点并计算距离
            query = point_a[torch.arange(B), center_point].unsqueeze(1)  # [B,1,3]
            dist = torch.norm(point_a - query, dim=2)  # [B,1024]

            # 选择最近的int_lam个点
            _, idxs = torch.topk(dist, int_lam, dim=1, largest=False)

            # 应用混合
            batch_indices = torch.arange(B, device=device)[:, None]
            # 确保 original_points 的形状是 [B, 1024, 3]
            assert original_points.shape[2] == 3, f"original_points tensor should have 3 channels (coordinates), but got {original_points.shape[2]}"
            original_points[batch_indices, idxs, :] = point_c[batch_indices, idxs, :]

            # 将修改后的 original_points 再进行维度交换
            points = original_points.transpose(2, 1)

            # 计算最终lambda
            lam = int_lam / 1024.0

            # 前向传播
            if 'lowfreq' in args.model:
                pred = model(points[:, :, :3], original_feats)  # 仅使用坐标
                trans_feat = None  # 若该模型没有 trans_feat 返回值，可设为 None
            else:
                pred, trans_feat = model(points)

            loss = criterion(pred, target_a) * (1. - lam) + criterion(pred, target_b) * lam
        else:
            if 'lowfreq' in args.model:
                pred = model(points[:, :, :3], original_feats)  # 仅使用坐标
                trans_feat = None  # 若该模型没有 trans_feat 返回值，可设为 None
            else:
                pred, trans_feat = model(points)
            loss = criterion(pred, target)

        if args.feature_transform and args.model == 'pointnet_kcutmix':
            loss += feature_transform_regularizer(trans_feat) * args.lambda_ft
        loss.backward()
        optimizer.step()

        pred_choice = pred.data.max(1)[1]
        correct += pred_choice.eq(target.data).cpu().sum()
        total += target.size(0)

        progress_bar(i, len(train_loader), 'Train Loss: %.3f | Train Acc: %.3f%% (%d/%d)'
                     % (loss.item() / (i + 1), 100. * correct.item() / total, correct, total))

    train_loss, train_acc = loss.item() / (i + 1), 100. * correct.item() / total

    ### Test in batch
    test_loss, test_acc = test(model, test_loader, criterion)
    acc_list.append(test_acc)
    print('the best test acc is', max(acc_list))

    ### Keep tracing
    log_row(logname, [epoch, train_loss, train_acc, test_loss, test_acc,
                      optimizer.param_groups[0]['lr'], max(acc_list), np.argmax(acc_list) - 1])
    save_ckpt(args, epoch, model, optimizer, acc_list)



